import ml_collections
import joblib


config = {
    # 保存预测结果的目录
    'save_path': 'data/results_4HFF',

    ###########################################################################################################
    # 单体模板和复合物参考结构
    ###########################################################################################################
    # 单体模板的pdb结构
    'template': 'data/4HFF.pdb',

    # [可选]复合物参考结构,用于计算RMSD,如果没有则设置为None.
    'native': './data/4HFF.pdb',

    # 对接的链号
    # 顺序决定对接时拼接的顺序.
    'chains': 'A,B',

    # 固定的链号,让指定链之间的相对位置与模板保持不变,没有要求则设置为None
    # 例如:
    #     'fixed_chains': ['A,B', 'C,D']
    #     A-B链的相对位置不变,C-D链的相对位置不变.
    #     'fixed_chains': ['A,B'],
    'fixed_chains': None,

    ###########################################################################################################
    # 实验约束
    ###########################################################################################################
    # 实验约束的距离约束,默认为8Å.
    # 由于AF2的定义限制,阈值需要在2-22Å之间
    'res_thres': 8.0,

    # 1v1 约束
    # 描述:
    #     距离在给定阈值以下的2个异链残基,没有则为None.
    #     多对1v1约束用[]进行罗列.
    #     残基序号指的是它在完整序列中的序号,序号从1开始.
    #         完整序列由所有指定的链的序号拼接,拼接顺序由`chains`定义.
    # 示例:
    #     'rest_1v1': [55,90]:整条链上的第55号残基和90号残基距离低于规定阈值
    #     'rest_1v1': [[78,198],[20,50]]:整条链上的第78号残基与198号残基距离/第20号残基与50号残基距离均低于规定阈值
    'rest_1v1': [79, 199],

    # 1vN 约束
    # 描述:
    #     距离在给定阈值以下的1个残基与另1组残基,没有则为None
    # 示例:
    #     'rest_1vN': [36,list(range(160,171))+[178,190]]
    #     整条链上36号残基与[160~170,178,190]号残基中至少有1个距离满足约束.注意range是左闭右开的
    'rest_1vN': None,

    # MvN 约束
    # 描述:
    #     包含若干个1vN约束,只有给定数量的残基对满足约束,没有则为None
    #     多个MvN约束用[]进行罗列
    # 示例:
    #     'rest_MvN': [[10, list(range(160, 170))],
    #                  [78, list(range(160, 170))],
    #                  [120, list(range(160, 170))],
    #                  2]
    #     3个给定的1vN约束,其中只有2个满足约束.
    # 也支持从pkl数据文件中导入(由extract_rest.py生成)
    # 'rest_MvN': joblib.load('./protein/1AHW/rest_MvN.pkl'),
    'rest_MvN': None,

    # 排斥距离阈值,两个残基C-beta原子的距离需要超过该距离,需要在2-22埃范围内
    'rep_thres': 8.0,

    # 排斥约束,与1v1约束规定类似,但距离要超过阈值
    # 示例:
    #     'rest_rep: [154, 250]
    #     整条链上第154,250号残基距离超过阈值
    'rest_rep': None,

    ###########################################################################################################
    # 优化参数
    ###########################################################################################################
    # 基于片段的优化切割长度,若没有则为None.
    # 对于长链序列,建议值为200
    # 可以减少显存消耗,但可能陷入局部最优化
    'crop_len': 200,

    # 重复预测的轮次数(rounds)
    # 增大数值会提高精度,但会增加模拟时间
    # 论文的案例重复5次
    'rounds': 1,

    # 每轮次反向传播的次数(steps),让结构收敛
    # 基于片段的优化建议使用更大的数值,如150(论文案例为100)
    # 否则建议设置为50
    'steps': 100,

    # 每n次反向传播保存1次构象.
    # 对基于片段的优化很有用,以缩减生成结构的文件大小,建议设置为3以上
    # 普通优化过程设置为1即可
    'save_every_n_step': 10,

    ###########################################################################################################
    # AF2 模型
    ###########################################################################################################
    # AF2 权重参数所在路径
    'data_dir': '/home/regen/projects/af_para',
    # 使用AF2-multimer,设为False则使用AF2
    'use_multimer': True,

    # 是否使用bfloat浮点数精度(速度上升,精度下降)
    'bfloat': True,
}

config = ml_collections.ConfigDict(config)
